#!/usr/bin/env python 
# -*- coding: utf-8 -*-
# @Time    : 2019/1/18 16:44
# @Author  : Tang Yang
# @Desc    : 
# @File    : result_analysis.py
from openpyxl import Workbook
import numpy as np
import json

TITLE = ["照片编号", "品类", "标记代码", "样本总数", "人工标记数", "识别数", "准确框", "标签错误框", "多框", "漏框", "错误识别的类别"]
ANALYSIS_KEYS = ["image_name", "category", "label", "sample_nums", "gt_nums", "recog_nums",
                 "correct_nums", "label_error_nums", "extra_nums", "miss_nums", "error_labels"]


def __is_overlap(rect1, rect2):
    """
    判断矩形是否存在重叠

    Parameters
    ----------
    rect1: 矩形1，形如[axis1_min, axis2_min, axis1_max, axis2_max]
    rect2: 矩形2，形如[axis1_min, axis2_min, axis1_max, axis2_max]

    Returns
    -------

    """
    return not ((rect1[0] >= rect2[2]) or
                (rect1[1] >= rect2[3]) or
                (rect1[2] <= rect2[0]) or
                (rect1[3] <= rect2[1]))


def __get_overlap_area(rect1, rect2):
    """
    计算两个矩形的重叠部分的面积，如果不重叠返回0

    Parameters
    ----------
    rect1: 矩形1，形如[axis1_min, axis2_min, axis1_max, axis2_max]
    rect2: 矩形2，形如[axis1_min, axis2_min, axis1_max, axis2_max]

    Returns
    -------

    """
    if not __is_overlap(rect1, rect2):
        return 0.0
    xmin = max(rect1[0], rect2[0])
    ymin = max(rect1[1], rect2[1])
    xmax = min(rect1[2], rect2[2])
    ymax = min(rect1[3], rect2[3])
    return (xmax - xmin) * (ymax - ymin)


def _eval_detect_result(true_boxes, true_classes, pred_boxes, pred_classes,
                        threshold=0.3, default_class=0):
    ground_true = []
    predictions = []
    _index = []
    for box, cls in zip(pred_boxes, pred_classes):
        overlap_rate = 0.0
        response_cls = -1
        for idx, (t_box, t_cls) in enumerate(zip(true_boxes, true_classes)):
            if __is_overlap(t_box, box):
                area = __get_overlap_area(t_box, box)
                src_area = min((box[2] - box[0]) * (box[3] - box[1]),
                               (t_box[2] - t_box[0]) * (t_box[3] - t_box[1]))
                if (area / src_area) > threshold:
                    if (area / src_area) > overlap_rate:
                        overlap_rate = area / src_area
                        response_cls = t_cls
                        _index.append(idx)
        ground_true.append(default_class if response_cls == -1 else response_cls)
        predictions.append(cls)
    for idx in range(len(true_classes)):
        if idx not in _index:
            ground_true.append(true_classes[idx])
            predictions.append(default_class)
    return ground_true, predictions


class ResultAnalyzer:
    """
    对检测结果进行分析，并且提供将结果写入Excel的功能
    """

    def __init__(self, workbook_save_path, all_sample_nums_path):
        with open(all_sample_nums_path, "r") as f:
            self._all_sample_nums = json.load(f)
        self._wb = Workbook()
        self._save_path = workbook_save_path
        self._ws = self._wb.create_sheet()
        self._start_row = 1
        for col, content in enumerate(TITLE):
            self._ws.cell(row=self._start_row, column=col + 1, value=content)
        self._start_row += 1
        self._analysis_result = []
        self._c_m = None
        self._all_gt = []
        self._all_pred = []

    def analysis_result(self, image_name, category, label_map_dict, gt_boxes, gt_classes, prediction_boxes,
                        prediction_classes):
        """
        对一张检测结果进行分析，返回分析结果
        :param prediction_classes:
        :param prediction_boxes:
        :param gt_classes:
        :param gt_boxes:
        :param label_map_dict:
        :param category: 是属于哪个品类，比如caffe、maipian...
        :param image_name:
        :return: 形如 [{"image_name": "1.jpg", "category": "caffe", 'label': "nco1", "sample_nums": 321, "gt_nums": 6,
                        recog_nums: 10, correct_nums: 8, label_error_nums: 2, extra_nums: 2, miss_nums: 4}, ... ] 的结果
                  列表中一行表示这张图片中一个类别的分析结果
        """
        gt, prediction = _eval_detect_result(gt_boxes, gt_classes, prediction_boxes, prediction_classes)
        self._all_gt += gt
        self._all_pred += prediction
        assert len(gt) == len(prediction), "length of gt must equal to length of prediction"
        dict_ret = {"image_name": image_name, "category": category}
        label_index_set = (set(gt) | set(prediction))
        if 0 in label_index_set:
            label_index_set.remove(0)
        for label_index in label_index_set:
            temp_dict = dict_ret.copy()
            temp_dict["label"] = label_map_dict[int(label_index)]
            if temp_dict["label"] not in self._all_sample_nums:
                sample_nums = -1
            else:
                sample_nums = self._all_sample_nums[temp_dict["label"]]  # 此类别的样本总数
            gt_nums = 0  # 此类别的人工标记数
            recog_nums = 0  # 此类别的模型检测数
            correct_nums = 0  # 此类别模型检测数中正确的个数
            label_error_nums = 0  # 此类别模型检测数中标签错误的个数
            extra_nums = 0  # 此类别模型检测数中多检的个数
            miss_nums = 0  # 漏检的个数
            error_labels = {}
            for i in range(len(gt)):
                gt_nums += 1 if gt[i] == int(label_index) else 0
                recog_nums += 1 if prediction[i] == int(label_index) else 0
                correct_nums += 1 if prediction[i] == int(label_index) and gt[i] == int(label_index) else 0
                label_error_nums += 1 if prediction[i] == int(label_index) and prediction[i] != gt[i] and gt[
                    i] != 0 else 0
                extra_nums += 1 if prediction[i] == int(label_index) and gt[i] == 0 else 0
                miss_nums += 1 if gt[i] == int(label_index) and prediction[i] == 0 else 0
                if prediction[i] == int(label_index) and gt[i] != int(label_index) and gt[i] != 0:
                    label = label_map_dict[int(gt[i])]
                    if label not in error_labels:
                        error_labels[label] = 1
                    else:
                        error_labels[label] += 1
            temp_dict.update(
                dict(sample_nums=sample_nums, gt_nums=gt_nums, recog_nums=recog_nums, correct_nums=correct_nums,
                     label_error_nums=label_error_nums, extra_nums=extra_nums, miss_nums=miss_nums,
                     error_labels=error_labels))
            self._analysis_result.append(temp_dict)
        return self._analysis_result

    def write_single_image_result_2_sheet(self, single_image_result=None):
        analysis_result = self._analysis_result if single_image_result is None else single_image_result
        if analysis_result is None:
            raise ValueError("analysis_result() must be called first.")
        for idx in range(len(analysis_result)):
            for col, key in enumerate(ANALYSIS_KEYS):
                if key == "error_labels":
                    self._ws.cell(row=self._start_row + idx, column=col + 1,
                                  value=str(analysis_result[idx][key])[1:-1])
                else:
                    self._ws.cell(row=self._start_row + idx, column=col + 1, value=analysis_result[idx][key])
        self._start_row += 1

    def analysis_and_write_single_image_result_2_sheet(self, image_name, category, label_map_dict, gt_boxes, gt_classes,
                                                       prediction_boxes, prediction_classes):
        ar = self.analysis_result(image_name, category, label_map_dict, gt_boxes, gt_classes, prediction_boxes,
                                  prediction_classes)
        self.write_single_image_result_2_sheet(ar)

    def save(self, save_path=None):
        if save_path is None:
            save_path = self._save_path
        self._wb.save(save_path)


def main():
    pass


if __name__ == '__main__':
    main()
